#!/usr/bin/perl

use Scalar::Util qw(looks_like_number);

# Note: 
#  all rows from "non" key body sites will be removed.
#
#  all rows from the completed list will be removed. 
#
# Categories:
# STOOL: All stool samples from all 300 subjects
#
# 3PEAT: All 6 key body sites from all three visits of the 3-peat subjects
# 
# At least 1 visit from each of the 300 subjects will be sequenced -
# ideally, all 6 key body sites will be completed, but if the
# quantities are insufficient, then some key body sites will not be
# pursued. The focus should be on selecting the best visit and
# sequencing as many of the six body sites as possible for that visit.
# For this, there are two types:
#   BEST_QUANTITY
#   BEST_VOLUME
# 
# 2PEAT: Samples the 2-peat subjects will be used to complete the resource.
# 
# These categories represent a set of rules, and for each rule remove
#   the sample from processing from the next consecutive rule. 

use strict;

my $list;
my $count = 0;

my $MASS_CUTOFF = 40;
my $VOLUME_CUTOFF = 10;

my @cols = ("barcode", "current_center", "sampling_center", "rsid", "gender", "specimen_type", "visit_num", "volume", "ctr_volume", "ng_per_ul", "in_library_construction", "ii_seq_june_2012");

my $body_site_column = 5;
my $barcode_column = 0;

my @key_sites = ("Stool", "Tongue dorsum", "Supragingival plaque", "Buccal mucosa", "Anterior nares", "Posterior fornix");

my $exclude = load_list(@ARGV[0]);
my $quants = load_quants(@ARGV[1]);

while(<STDIN>) {
  s/\n//;
  my $i = 0;

  my @data = split(/\t/);
  my $site = $data[$body_site_column];
  my $rsid = $data[$barcode_column];

  foreach my $s (@key_sites) {
      if ($site =~ /$s/ && $exclude->{$rsid} != 1) {
	  foreach my $j (@data) {
	      $list->{$count}->{$cols[$i]} = $j;
	      $i++
	  }
	  # changes tag to one in body site list
	  $list->{$count}->{'specimen_type'} = $s; 

	  # add in quantities supplied from ARGV[1]
	  # to replace fields containing ctr_volume ng_per_ul
	  if (length($quants->{$rsid}->{'volume'} != 0)) {
	      $list->{$count}->{'ctr_volume'} = $quants->{$rsid}->{'volume'};
	  }
	  if (length($quants->{$rsid}->{'quant'} != 0)) {
	      $list->{$count}->{'ng_per_ul'} = $quants->{$rsid}->{'quant'};
	  }

	  $list->{$count}->{'STOOL'} = 0;
	  $list->{$count}->{'3PEAT'} = 0;
	  $list->{$count}->{'BEST_QUANTITY'} = 0;
	  $list->{$count}->{'BEST_VOLUME'} = 0;
	  $list->{$count}->{'2PEAT'} = 0;
	  $count++;
      }
  }
}

$list->{'count'} = $count;

assign_mass($list);
find_stool($list);
find_three_peat($list);
find_best_quantity($list, @key_sites);
find_best_volume($list, @key_sites);
find_two_peat($list);

my $i;

print "barcode\tSUBJECTid\tgender\tcurrent_center\tspecimen_type\tvisit_num\tctr_volume\tng_per_ul\tmass\tstool\t3peat\tbest_quantity\tbest_volume\t2peat\tall_tags\n";

for($i=0;$i<$list->{'count'};$i++) {

    print "$list->{$i}->{'barcode'}\t$list->{$i}->{'rsid'}\t$list->{$i}->{'gender'}\t$list->{$i}->{'current_center'}\t$list->{$i}->{'specimen_type'}\t$list->{$i}->{'visit_num'}\t$list->{$i}->{'ctr_volume'}\t$list->{$i}->{'ng_per_ul'}\t$list->{$i}->{'mass'}\t";

    my $done = 0;
    if ($list->{$i}->{'STOOL'} == 1 && $done !=1) {
	print "STOOL\t";
	$done = 1;
    }
    else {
	print " \t";
    }
    if ($list->{$i}->{'3PEAT'} == 1 && $done !=1) {
	print "3PEAT\t";
	$done = 1;
    }
    else {
	print " \t";
    }
    if ($list->{$i}->{'BEST_QUANTITY'} == 1 && $done !=1) {
	print "BEST_QUANTITY\t";
	$done = 1;
    }
    else {
	print " \t";
    }
    if ($list->{$i}->{'BEST_VOLUME'} == 1 && $done !=1) {
	print "BEST_VOLUME\t";
	$done = 1;
    }
    else {
	print " \t";
    }
    if ($list->{$i}->{'2PEAT'} == 1 && $done !=1) {
	print "2PEAT\t";
	$done = 1;
    }
    else {
	print " \t";
    }

    $list->{$i}->{'tag'} =~ s/^:://;
    print "$list->{$i}->{'tag'}\n";
}

sub assign_mass {
    my $list = shift;
    for($i=0;$i<$list->{'count'};$i++) {
	my $mass;
	if ($list->{$i}->{'ctr_volume'} eq "\#N/A") {
	    $list->{$i}->{'ctr_volume'} = "N/A";
	}
	if (looks_like_number($list->{$i}->{'ctr_volume'}) && 
	    looks_like_number($list->{$i}->{'ng_per_ul'})) {
	    $mass = $list->{$i}->{'ctr_volume'} * 
		$list->{$i}->{'ng_per_ul'};
	}
	else {
	    $mass = "N/N";
	}
	$list->{$i}->{'mass'} = $mass;
    }
}

sub find_best_volume {
    my $list = shift;
    my @sites = @_;

    # get all subjects:
    my @subjects;
    my $visit_data;
    my $rsid_visit;
    my $done;
    my $i;
    # collect the volumes, a unique list of subjects
    #  and collect the rows that have the same visits from the patient
    for($i=0;$i<$list->{'count'};$i++) {
	my $rsid = $list->{$i}->{'rsid'};
	my $visit = $list->{$i}->{'visit_num'};
	my $site = $list->{$i}->{'specimen_type'};

	my $volume;
	if (looks_like_number($list->{$i}->{'ctr_volume'})) {
	    $volume = $list->{$i}->{'ctr_volume'}
	}
	else {
	    $volume = "N/N";
	}

	$visit_data->{$rsid}->{$visit}->{$site} = $volume;
	$rsid_visit->{$rsid}->{$visit} = $i;

	push(@subjects, $rsid) if ($done->{$rsid} != 1);
	$done->{$rsid} = 1;
    }

    my $best_visit;
    my $hit;
    # for each patient:
    # count how many good visits for each body site there were. 
    # and pick the best visit
    foreach my $rsid (@subjects) {
	my $max_count = 0;
	$best_visit->{$rsid} = "none";
	foreach my $visit (1,2,3) {
	    my $count = 0;
	    foreach my $site (@sites) {
		my $volume = $visit_data->{$rsid}->{$visit}->{$site};
		# print "RSID: $rsid :: $visit :: $site :: $volume\n";
		$count++ if (looks_like_number($volume) && $volume > $VOLUME_CUTOFF);
	    }
	    # print "  COUNT: $count\n";
	    if ($max_count < $count) {
		$max_count = $count;
		$best_visit->{$rsid} = $visit;
	    }
	}
	my $c = $rsid_visit->{$rsid}->{$best_visit->{$rsid}};
	if (looks_like_number($best_visit->{$rsid})) {
	    # print "  BEST: $rsid :: $best_visit->{$rsid}\n";
	    # collect the rsid and good visit
	    $hit->{$rsid}->{$best_visit->{$rsid}} = 1;
	}
	else {
	    $list->{$c}->{'fail_visit'} = 1;
	    # print "  Not assigned\n";
	}
    }

    for($i=0;$i<$list->{'count'};$i++) {
	my $rsid = $list->{$i}->{'rsid'};
	my $visit = $list->{$i}->{'visit_num'};
	if ($hit->{$rsid}->{$visit} == 1) {
	    $list->{$i}->{'BEST_VOLUME'} = 1;
	    $list->{$i}->{'tag'} = $list->{$i}->{'tag'} . "::" . "BEST_VOLUME";
	}
    }
}

sub find_best_quantity {
    my $list = shift;
    my @sites = @_;

    # get all subjects:
    my @subjects;
    my $visit_data;
    my $rsid_visit;
    my $done;
    my $i;
    for($i=0;$i<$list->{'count'};$i++) {
	my $rsid = $list->{$i}->{'rsid'};
	my $visit = $list->{$i}->{'visit_num'};
	my $site = $list->{$i}->{'specimen_type'};

	my $mass = $list->{$i}->{'mass'};
	    
	$visit_data->{$rsid}->{$visit}->{$site} = $mass;
	$rsid_visit->{$rsid}->{$visit} = $i;

	push(@subjects, $rsid) if ($done->{$rsid} != 1);
	$done->{$rsid} = 1;
    }

    my $best_visit;
    my $hit;
    # count how many good visits for each body site there were. 
    foreach my $rsid (@subjects) {
	my $max_count = 0;
	$best_visit->{$rsid} = "none";
	foreach my $visit (1,2,3) {
	    my $count = 0;
	    foreach my $site (@sites) {
		my $mass = $visit_data->{$rsid}->{$visit}->{$site};
		# print "RSID: $rsid :: $visit :: $site :: $mass\n";
		$count++ if (looks_like_number($mass) && $mass > $MASS_CUTOFF);
	    }
	    # print "  COUNT: $count\n";
	    if ($max_count < $count) {
		$max_count = $count;
		$best_visit->{$rsid} = $visit;
	    }
	}
	my $c = $rsid_visit->{$rsid}->{$best_visit->{$rsid}};
	if (looks_like_number($best_visit->{$rsid})) {
	    # print "  BEST: $rsid :: $best_visit->{$rsid}\n";
	    $hit->{$rsid}->{$best_visit->{$rsid}} = 1;
	}
	else {
	    $list->{$c}->{'fail_visit'} = 1;
	    # print "  Not assigned\n";
	}
    }

    for($i=0;$i<$list->{'count'};$i++) {
	my $rsid = $list->{$i}->{'rsid'};
	my $visit = $list->{$i}->{'visit_num'};
	if ($hit->{$rsid}->{$visit} == 1) {
	    $list->{$i}->{'BEST_QUANTITY'} = 1;
	    $list->{$i}->{'tag'} = $list->{$i}->{'tag'} . "::" . "BEST_QUANTITY";
	}
    }
}

sub find_stool {
    my $list = shift;

    my $i;
    for($i=0;$i<$list->{'count'};$i++) {
	if ($list->{$i}->{'specimen_type'} =~ /Stool/) {
	    # print "Rule 1: $i\t $list->{$i}->{'specimen_type'}\n";
	    $list->{$i}->{'STOOL'} = 1;
	    $list->{$i}->{'tag'} = $list->{$i}->{'tag'} . "::" . "STOOL";
	}
    }
}

sub find_three_peat {
    my $list = shift;

    my $i;
    my $hit;
    for($i=0;$i<$list->{'count'};$i++) {
	if ($list->{$i}->{'visit_num'} == 3) {
	    # save the rsid and visit
	    $hit->{$list->{$i}->{'rsid'}} = 1;
	}
    }

    for($i=0;$i<$list->{'count'};$i++) {
	if ($hit->{$list->{$i}->{'rsid'}} == 1 &&
	    $list->{$i}->{'visit_num'} == 3) {

	    $list->{$i}->{'3PEAT'} = 1;
	    $list->{$i}->{'tag'} = $list->{$i}->{'tag'} . "::" . "3PEAT";
	}
    }
}

sub find_two_peat {
    my $list = shift;

    my $i;
    my $hit;
    for($i=0;$i<$list->{'count'};$i++) {
	if ($list->{$i}->{'visit_num'} == 2) {
	    # save the rsid and visit
	    $hit->{$list->{$i}->{'rsid'}} = 1;
	}
    }

    for($i=0;$i<$list->{'count'};$i++) {
	if ($hit->{$list->{$i}->{'rsid'}} == 1 &&
	    $list->{$i}->{'visit_num'} == 2) {
	    $list->{$i}->{'2PEAT'} = 1;
	    $list->{$i}->{'tag'} = $list->{$i}->{'tag'} . "::" . "2PEAT";
	}
    }
}

sub load_list {
    my $file = shift;
    my $l;

    open(F_IN, $file) || die ("supply an exlude list");
    while(<F_IN>) {
	s/\n//;
	$l->{$_} = 1;
    }

    return ($l)
}

sub load_quants {
    my $file = shift;
    my $l;

    # The quants files are of the format:
    # ID VOLUME QUANT

    open(F_IN, $file) || die ("supply an quant list");
    while(<F_IN>) {
	s/\n//;
	my ($id, $vol, $q) = split("\t");
	$l->{$id}->{'volume'} = $vol;
	$l->{$id}->{'quant'} = $q;
    }

    return ($l)
}

